package com.skw.nio;

import java.io.IOException;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.Socket;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.nio.channels.spi.SelectorProvider;
import java.util.*;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

public class EchoServer {
    private Selector selector;
    private ExecutorService tp;
    public static Map<Socket, Long> time_stat;

    public EchoServer() {
        tp = Executors.newCachedThreadPool();
        time_stat = new HashMap<Socket, Long>(10240);
    }

    private void startServer() throws IOException {
        selector = SelectorProvider.provider().openSelector();

        ServerSocketChannel ssc = ServerSocketChannel.open();
        //设置为非阻塞
        ssc.configureBlocking(false);

        // 进行端口绑定，将这个（ServerSocket）Channel绑定在8000端口上
        //InetSocketAddress isa = new InetSocketAddress(InetAddress.getLocalHost(), 8000);
        InetSocketAddress isa = new InetSocketAddress(8000);
        ssc.socket().bind(isa);

        // 将这个ServerSocketChannel绑定到Selector上，并注册它的感兴趣时间为Accept
        SelectionKey acceptKey = ssc.register(selector, SelectionKey.OP_ACCEPT);

        for (;;) {
            // 阻塞方法，返回值是准备就绪的SelectionKey的数量
            selector.select();
            // 获取那些准备好的SelectionKey
            Set readyKeys = selector.selectedKeys();
            Iterator i = readyKeys.iterator();
            long e = 0;
            while (i.hasNext()) {
                SelectionKey sk = (SelectionKey)i.next();
                // 重要步骤，不移除则会重复处理相同的SelectionKey
                i.remove();
                // Channel 是否在Accept不了状态
                if (sk.isAcceptable()) {
                    doAccept(sk);
                }// Channel 是否在可读状态
                else if (sk.isValid() && sk.isReadable()) {
                    if (!time_stat.containsKey(((SocketChannel)sk.channel()).socket())) {
                        time_stat.put(((SocketChannel)sk.channel()).socket(), System.currentTimeMillis());
                    }
                    doRead(sk);
                } // Channel是否可写状态
                else if (sk.isValid() && sk.isWritable()) {
                    doWrite(sk);
                    e = System.currentTimeMillis();
                    long b = time_stat.remove(((SocketChannel)sk.channel()).socket());
                    System.out.println("spend:" + (e-b) + "ms");
                }
            }
        }
    }

    private void doAccept(SelectionKey sk) {
        ServerSocketChannel server = (ServerSocketChannel) sk.channel();
        //当有一个新客户端连接接入时，就会产生一个新的Channel来代表这个连接
        SocketChannel clientChannel;
        try {
            clientChannel = server.accept();
            // Channel配置为非阻塞式，系统在准备好IO后再通知（Selector）线程来读取或写入
            clientChannel.configureBlocking(false);

            // 将生成的Channel注册到Selector上，并告诉Selector现在对读感兴趣
            SelectionKey clientKey = clientChannel.register(selector, SelectionKey.OP_READ);
            // 创建一个EchoClient客户端示例
            ClientStorage echoClient = new ClientStorage();
            // 客户端附加到表示这个连接的SelectionKey上
            clientKey.attach(echoClient);

            InetAddress clientAddress = clientChannel.socket().getInetAddress();
            System.out.println("Accept connection from " + clientAddress.getHostAddress() + ".");
        } catch (IOException e) {
            System.out.println("Falied to accept new client.");
            e.printStackTrace();
        }
    }

    //注：读和写使用同一个SelectionKey，读写操作共享一个ClientStorage示例
    public void doRead(SelectionKey sk) {
        // 得到当前客户端的Channel
        SocketChannel channel = (SocketChannel)sk.channel();
        // 准备8k的缓冲区读取数据
        ByteBuffer bb = ByteBuffer.allocate(8192);
        int len;

        try {
            len = channel.read(bb);
            if (len < 0) {
                disconnect(sk);
                return;
            }
        } catch (IOException e) {
            System.out.println("Failed to read from client");
            e.printStackTrace();
            disconnect(sk);
            return;
        }
        // 读取完成后，重置缓冲区，指针复位，头指针指向头，尾指针指向limit，为数据处理做准备
        bb.flip();
        // 为模拟更复杂场景，不阻塞任务派发线程，使用线程池
        tp.execute(new HandleMsg(sk, bb));
    }

    public void doWrite(SelectionKey sk) {
        SocketChannel channel = (SocketChannel)sk.channel();
        ClientStorage echoClient = (ClientStorage)sk.attachment();
        LinkedList<ByteBuffer> outq = echoClient.getOutputQueue();

        ByteBuffer bb = outq.getLast();
        try {// 写回操作
            int len = channel.write(bb);
            if (len == -1) {
                disconnect(sk);
                return;
            }
            // 如果bb全部发送完成，则移除这个缓存对象
            if (bb.remaining() == 0) {
                outq.removeLast();
            }
        } catch (IOException e) {
            System.out.println("Failed to write to client.");
            e.printStackTrace();
            disconnect(sk);
        }
        //全部数据发送完成后（也就是outq为0），需要将写时间（OP_WRITE）从感兴趣的操作中移除
        if (outq.size() == 0) {
            sk.interestOps(SelectionKey.OP_READ);
        }
    }

    private void disconnect(SelectionKey sk) {
        SocketChannel channel = (SocketChannel)sk.channel();

        InetAddress clientAddress = channel.socket().getInetAddress();
        System.out.println(clientAddress.getHostAddress());
        try {
            channel.close();
        } catch (IOException e) {
            System.out.println("Falied to close client socket channel.");
            e.printStackTrace();
        }
    }

    class HandleMsg implements Runnable {
        SelectionKey sk;
        ByteBuffer bb;

        public HandleMsg(SelectionKey sk, ByteBuffer bb) {
            this.sk = sk;
            this.bb = bb;
        }
        @Override
        public void run() {
            ClientStorage echoClient = (ClientStorage) sk.attachment();
            echoClient.enqueue(bb);
            //在数据处理完成后，就准备写回到客户端
            //因此，重新注册感兴趣的消息时间，将写入操作也作为感兴趣的时间提交
            sk.interestOps(SelectionKey.OP_READ | SelectionKey.OP_WRITE);
            // 强迫Selector立即返回
            selector.wakeup();
        }
    }

    public static void main(String[] args) {
        EchoServer echoServer = new EchoServer();
        try {
            echoServer.startServer();
        } catch (Exception e) {
            System.out.println("Exception caught, program exiting...");
            e.printStackTrace();
        }
    }
}
